from business.UnpackedData import UnpackedData 
from business.HuffmanTree import HuffmanTree, createHuffmanTreeFromSymbolDict

import logging

class DecryptFile:
    def __init__(self):
        self.mHeaderDataDef = UnpackedData("DecryptFileHeader",
                                           [["id"                  , "4s"],
                                            ["uncompressedFileSize", "I"]])
        self.mHeader = None
        self.mInputData = None
        self.mOuputData = None
        
        self.mSymbolRepartitionHuffmanTree = fillSymbolRepartitionHuffmanTree()
        
    def decrypt(self, iFile, iFileOffset, iSize):   
        logging.info("DecryptFile: Offset(" + str(iFileOffset) + ") Size(" + str(iSize) + ")")
        
        iFile.seek(iFileOffset)
        lInput = iFile.read(iSize)
        lCurrentPos = 0
        self.mHeader, lCurrentPos = self.mHeaderDataDef.unpackWithOffset(lInput, lCurrentPos)
        logging.debug(self.mHeader)
        
        lBitList = list()
        for lIndex in range(lCurrentPos // 4, len(lInput) // 4): # We go through the bytes in uint32 big endian/network order
            if (lIndex + 1) % (1 << 14) != 0: # We skip the last bytes for each 65536 bytes chunk, maybe a crc ?
                for lByteNo in range(3, -1, -1):
                    lCurrentByte = lInput[4 * lIndex + lByteNo]
                    lBitList.append(bin(lCurrentByte).lstrip('0b').rjust(8, '0'))
        lBitString = "".join(lBitList)
        
        lPos = 0
        lConstValueWriteSizeAddition = int(lBitString[lPos+4:lPos + 8], 2) + 1
        logging.debug("lConstValueWriteSizeAddition: " + str(lConstValueWriteSizeAddition))
        lPos += 8
        
        lOutput = bytearray()
        try:
            while (len(lOutput) < self.mHeader.uncompressedFileSize):
                logging.debug("Processing HuffmanTreeCode...")
                lPos, lHuffmanTreeCode = self._parseHuffmanTree(lBitString, lPos)
                logging.debug("Processing HuffmanTreeCopy...")
                lPos, lHuffmanTreeCopy = self._parseHuffmanTree(lBitString, lPos)
                
                lMaxCount = (int(lBitString[lPos:lPos + 4], 2) + 1) << 12      
                logging.debug("lMaxCount: " + str(lMaxCount))            
                lPos += 4
                
                lCurrentCout = 0
                while ((lCurrentCout < lMaxCount) and (len(lOutput) < self.mHeader.uncompressedFileSize)):
                    logging.debug("lCurrentCount: " + str(lCurrentCout) + " - lPos: " + str(lPos))
                    lCurrentCout += 1
                    lCode, lOffset = lHuffmanTreeCode.readCode(lBitString, lPos)
                    lPos += lOffset
                
                    if (lCode <= 0xFF): 
                        lOutput.append(lCode)
                    else: # more than one byte, it means we are in copy mode
                        lCode -= 0x100 
                        lDiv, lMod = divmod(lCode, 4)
                        lWriteSize = lCode
                        if lDiv > 0:
                            lWriteSize = (1 << (lDiv - 1)) * (4 + lMod) 
                            if lCode == 28:
                                lWriteSize = 0xFF # special case for 28 don't need more bits as we already have 0xFF/seems to be the highest value possible
                            elif lCode > 28:
                                logging.error("Should never happen lCode > 28")
                                lWriteSize = 0
                            if 1 < lDiv < 7:
                                lWriteSize |= (int(lBitString[lPos:lPos + lDiv - 1], 2)) # reading lDiv - 1 extra bits/number of bits that were lshift
                                lPos += lDiv - 1
                        lWriteSize += lConstValueWriteSizeAddition
                        
                        # Reading the code for the Write Location
                        lCode, lOffset = lHuffmanTreeCopy.readCode(lBitString, lPos)
                        lPos += lOffset
                        lDiv, lMod = divmod(lCode, 2)
                        lWriteOffset = lCode
                        if lDiv > 0:
                            lWriteOffset = ((1 << (lDiv - 1)) * (2 + lMod))
                            if lDiv > 1:
                                lWriteOffset |= (int(lBitString[lPos:lPos + lDiv - 1], 2))
                                lPos += lDiv - 1
                        lWriteOffset += 1
                        
                        logging.debug("Copy mode, offset: " + str(lWriteOffset) + " - size: " + str(lWriteSize))
                        
                        if (lWriteOffset > len(lOutput)):
                            logging.error("WriteOffset > len(lOutput)")
                            raise Exception from "WriteOffset > len(lOutput)"
                        else:
                            lWriteLoc = len(lOutput) - lWriteOffset
                            lWriteAvailable = lWriteOffset
                            while (lWriteSize > 0):
                                if lWriteSize > lWriteAvailable:
                                    lOutput.extend(lOutput[lWriteLoc:lWriteLoc + lWriteAvailable])
                                    lWriteSize -= lWriteAvailable
                                else:
                                    lOutput.extend(lOutput[lWriteLoc:lWriteLoc + lWriteSize])
                                    lWriteSize = 0
       
        except Exception as ex:
            logging.error("Exception catched while trying to decode the file: " + str(ex))
            return lOutput
            
        return lOutput
        
    def _parseHuffmanTree(self, iBitString, iOffset):
        lTotalNumberOfsymbols = 256 * int(iBitString[iOffset:iOffset+8], 2) + int(iBitString[iOffset+8:iOffset+16], 2)
        iOffset += 16
        lRemainingSymbols = lTotalNumberOfsymbols
        lBitsClassifiedCodes = dict()
        while (lRemainingSymbols > 0):
            lCode, lOffset = self.mSymbolRepartitionHuffmanTree.readCode(iBitString, iOffset)
            iOffset += lOffset
            if lCode.mNbOfBits != 0:
                if not(lCode.mNbOfBits in lBitsClassifiedCodes):
                    lBitsClassifiedCodes[lCode.mNbOfBits] = list()
                for lIndex in range(lCode.mNbOfSymbols):
                    lBitsClassifiedCodes[lCode.mNbOfBits].append(lRemainingSymbols - lIndex - 1)
            lRemainingSymbols -= lCode.mNbOfSymbols
        lHuffmanTree = createHuffmanTreeFromSymbolDict(lBitsClassifiedCodes)
        return iOffset, lHuffmanTree
    
class SymbolRepartitionHuffmanTreeValue:
    def __init__(self, iHexValue):
        self.mNbOfBits = iHexValue & 0x1F
        self.mNbOfSymbols = (iHexValue >> 5) + 1
    
    def __str__(self):
        return str(vars(self))

def fillSymbolRepartitionHuffmanTree():
    lHuffmanTree = HuffmanTree()
    
    lBitsClassifiedCodes = dict()
    lBitsClassifiedCodes[3] = [ 0x08, 0x09, 0x0A ]
    lBitsClassifiedCodes[4] = [ 0x00, 0x07, 0x0B, 0x0C ]
    lBitsClassifiedCodes[5] = [ 0x06, 0x29, 0x2A, 0xE0 ]
    lBitsClassifiedCodes[6] = [ 0x04, 0x05, 0x20, 0x28, 0x2B, 0x2C, 0x40, 0x4A ]
    lBitsClassifiedCodes[7] = [ 0x03, 0x0D, 0x25, 0x26, 0x27, 0x48, 0x49 ]
    lBitsClassifiedCodes[8] = [ 0x24, 0x47, 0x4B, 0x4C, 0x69, 0x6A ]
    lBitsClassifiedCodes[9] = [ 0x23, 0x46, 0x60, 0x63, 0x67, 0x68, 0x88, 0x89, 
                                0xA0, 0xE8 ]
    lBitsClassifiedCodes[10] = [ 0x01, 0x02, 0x2D, 0x43, 0x44, 0x45, 0x65, 0x66, 
                                 0x80, 0x87, 0x8A, 0xA8, 0xA9, 0xC0, 0xC9, 0xE9 ]
    lBitsClassifiedCodes[11] = [ 0x0E, 0x4D, 0x64, 0x6B, 0x6C, 0x84, 0x85, 0x8B, 
                                 0xA4, 0xA5, 0xAA, 0xC8, 0xE5 ]
    lBitsClassifiedCodes[12] = [ 0x83, 0x86, 0xA6, 0xA7, 0xC7, 0xCA, 0xE7 ]
    lBitsClassifiedCodes[13] = [ 0x22, 0x2E, 0x8C, 0xC4, 0xE4, 0xE6 ]
    lBitsClassifiedCodes[14] = [ 0x4E, 0x6D, 0xC6, 0xEC ]
    lBitsClassifiedCodes[15] = [ 0x0F, 0x10, 0x11, 0x8D, 0xAB, 0xAC, 0xCC, 0xEA ]
    lBitsClassifiedCodes[16] = [ 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19,
                                 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F, 0x21, 0x2F,
                                 0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 
                                 0x38, 0x39, 0x3A, 0x3B, 0x3C, 0x3D, 0x3E, 0x3F,
                                 0x41, 0x42, 0x4F, 0x50, 0x51, 0x52, 0x53, 0x54, 
                                 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x5B, 0x5C,
                                 0x5D, 0x5E, 0x5F, 0x61, 0x62, 0x6E, 0x6F, 0x70, 
                                 0x71, 0x72, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78,
                                 0x79, 0x7A, 0x7B, 0x7C, 0x7D, 0x7E, 0x7F, 0x81, 
                                 0x82, 0x8E, 0x8F, 0x90, 0x91, 0x92, 0x93, 0x94,
                                 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A, 0x9B, 0x9C, 
                                 0x9D, 0x9E, 0x9F, 0xA1, 0xA2, 0xA3, 0xAD, 0xAE,
                                 0xAF, 0xB0, 0xB1, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 
                                 0xB7, 0xB8, 0xB9, 0xBA, 0xBB, 0xBC, 0xBD, 0xBE,
                                 0xBF, 0xC1, 0xC2, 0xC3, 0xC5, 0xCB, 0xCD, 0xCE, 
                                 0xCF, 0xD0, 0xD1, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6,
                                 0xD7, 0xD8, 0xD9, 0xDA, 0xDB, 0xDC, 0xDD, 0xDE, 
                                 0xDF, 0xE1, 0xE2, 0xE3, 0xEB, 0xED, 0xEE, 0xEF,
                                 0xF0, 0xF1, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 
                                 0xF8, 0xF9, 0xFA, 0xFB, 0xFC, 0xFD, 0xFE, 0xFF ]
    
    for lKey in lBitsClassifiedCodes.keys():
        lBitsClassifiedCodes[lKey].reverse()
        for lIndex in range(len(lBitsClassifiedCodes[lKey])):
            lBitsClassifiedCodes[lKey][lIndex] = SymbolRepartitionHuffmanTreeValue(lBitsClassifiedCodes[lKey][lIndex])

    return createHuffmanTreeFromSymbolDict(lBitsClassifiedCodes)